use strum_macros::{AsRefStr, Display, EnumString};

use alloc::vec;

use crate::{
    dispatch_by_dtype,
    ir::tensor::{DType, Tensor},
    ops::{
        broadcast::{advance_coords, broadcast_aligned_strides, broadcast_shape},
        context::OpsContext,
        ops::Ops,
        ops_error::{OpsError, OpsResult},
    },
};

#[derive(Clone, Copy, Debug, PartialEq, Eq, AsRefStr, EnumString, Display)]
pub enum ElementwiseKind {
    Add,
    Sub,
    Mul,
    Div,
    Max,
    Min,
}

pub struct ElementwiseOps {
    pub kind: ElementwiseKind,
    pub exec_dtype: DType,
}

impl Ops for ElementwiseOps {
    fn name(&self) -> &str {
        self.kind.as_ref()
    }

    fn compute(&self, context: &mut crate::ops::context::OpsContext) -> OpsResult<()> {
        let dt = self.exec_dtype;
        let kind = self.kind;
        dispatch_by_dtype!(dt, |T| {
            match kind {
                ElementwiseKind::Add => run_binary::<T>(context, |a,b| a + b),
                ElementwiseKind::Sub => run_binary::<T>(context, |a,b| a - b),
                ElementwiseKind::Mul => run_binary::<T>(context, |a,b| a * b),
                ElementwiseKind::Div => run_binary::<T>(context, |a,b| a / b),
                ElementwiseKind::Max => run_binary::<T>(context, |a,b| if a > b { a } else { b }),
                ElementwiseKind::Min => run_binary::<T>(context, |a,b| if a < b { a } else { b }),
            }
        },
            DType::F32 => f32, DType::F64 => f64,
            DType::I8  => i8,  DType::I16 => i16, DType::I32 => i32, DType::I64 => i64,
            DType::U8  => u8,  DType::U16 => u16, DType::U32 => u32, DType::U64 => u64,
            // DType::Bool=> bool
        )
    }
}

#[inline]
pub fn run_binary<T: Copy>(context: &mut OpsContext, func: impl Fn(T, T) -> T) -> OpsResult<()> {
    if context.inputs.len() != 2 || context.outputs.is_empty() {
        return Err(OpsError::InvalidInput(
            "elementwise binary expects 2 inputs and 1 output",
        ));
    }

    let (input_a, input_b) = (&context.inputs[0], &context.inputs[1]);
    let output_tensor = &mut context.outputs[0];

    if input_a.dtype != output_tensor.dtype || input_b.dtype != output_tensor.dtype {
        return Err(OpsError::TypeMismatch(
            "dtype mismatch among inputs and outputs",
        ));
    }

    // shape inference (NumPy-style broadcasting)
    let output_shape = broadcast_shape(&input_a.shape, &input_b.shape)?;

    // rebuild output tensor if shape does not match
    if output_tensor.shape != output_shape {
        *output_tensor = Tensor::new(output_tensor.dtype, output_shape.clone());
    }

    // fast contiguous path
    if input_a.shape == output_shape && input_b.shape == output_shape {
        let input_a_slice = input_a.as_slice::<T>();
        let input_b_slice = input_b.as_slice::<T>();
        let output_slice = output_tensor.as_mut_slice::<T>();
        let element_count = output_slice.len();

        for i in 0..element_count {
            output_slice[i] = func(input_a_slice[i], input_b_slice[i]);
        }
        return Ok(());
    }

    // broadcasting path (row-major order + stride 0 for broadcasted dimensions)
    let output_rank = output_shape.len();
    let total_elements: usize = output_shape
        .iter()
        .fold(1usize, |acc, &dim| acc.saturating_mul(dim));

    // aligned strides after broadcasting
    let input_a_strides = broadcast_aligned_strides(&input_a.shape, &output_shape)?;
    let input_b_strides = broadcast_aligned_strides(&input_b.shape, &output_shape)?;

    let input_a_data = input_a.as_slice::<T>();
    let input_b_data = input_b.as_slice::<T>();
    let output_data = output_tensor.as_mut_slice::<T>();

    // coordinate odometer
    let mut coordinates = vec![0isize; output_rank];

    // source “linear index base”: since aligned strides are already in units
    // of element step size, coords·stride gives the linear offset directly.
    let mut linear_index = 0usize;
    loop {
        // compute source offsets
        let mut input_a_offset = 0isize;
        let mut input_b_offset = 0isize;
        for dim in 0..output_rank {
            input_a_offset += coordinates[dim] * input_a_strides[dim];
            input_b_offset += coordinates[dim] * input_b_strides[dim];
        }

        // write result
        output_data[linear_index] = func(
            input_a_data[input_a_offset as usize],
            input_b_data[input_b_offset as usize],
        );

        linear_index += 1;
        if !advance_coords(&mut coordinates, &output_shape) {
            break;
        }
    }

    debug_assert_eq!(linear_index, total_elements);
    Ok(())
}
